package p.ithorns.tools.translator.apt;

import com.sun.source.tree.MethodTree;
import com.sun.source.tree.Tree;
import com.sun.source.tree.VariableTree;
import com.sun.tools.javac.api.JavacTrees;
import com.sun.tools.javac.code.Types;
import com.sun.tools.javac.processing.JavacProcessingEnvironment;
import com.sun.tools.javac.tree.JCTree;
import com.sun.tools.javac.tree.TreeMaker;
import com.sun.tools.javac.util.Context;
import com.sun.tools.javac.util.List;
import com.sun.tools.javac.util.ListBuffer;
import com.sun.tools.javac.util.Names;
import org.reflections.Reflections;
import org.reflections.scanners.Scanners;
import org.reflections.util.ConfigurationBuilder;
import p.ithorns.tools.translator.apt.utils.AptUtil;
import p.ithorns.tools.translator.apt.utils.Parent;
import p.ithorns.tools.translator.apt.utils.Permit;
import p.ithorns.tools.translator.apt.utils.UnsafeUtil;
import sun.misc.Unsafe;

import com.sun.tools.javac.util.Name;
import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.Filer;
import javax.annotation.processing.Messager;
import javax.annotation.processing.ProcessingEnvironment;
import javax.lang.model.element.*;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.util.Elements;
import javax.tools.Diagnostic;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * BaseProcessor
 * <p>
 * 定义抽象的基类 BaseProcessor，用于统一获取编译阶段的工具类，如 JavacTrees、TreeMaker 等。
 *
 * @author Ht.L
 * @date 2024-05-30 23:18
 * @since 1.0.0
 */
@SuppressWarnings("unused")
public abstract class BaseProcessor extends AbstractProcessor {

    /**
     * 用来在编译期打log
     */
    protected Messager messager;

    /**
     * 提供了待处理的抽象语法树
     */
    protected JavacTrees trees;

    /**
     * 封装了创建AST节点的一些方法
     */
    protected TreeMaker maker;

    /**
     * 提供了创建标识符的方法
     */
    protected Names names;

    /**
     * 处理TypeMirror的工具类
     */
    protected Types types;

    /**
     * 处理Element的工具类
     */
    protected Elements elements;

    /**
     * 使用Filer可以创建文件
     */
    protected Filer filer;


    protected Context context;

    private final static String APT_HANDLER_PACKAGE = "p.ithorns.framework.translator.apt";

    @SuppressWarnings("MismatchedQueryAndUpdateOfCollection")
    private final static Map<String, AbstractHandler> MAP = new HashMap<>(16);


/**
     * Useful from jdk9 and up; required from jdk16 and up.
     * This code is supposed to gracefully do nothing on jdk8 and below, as this operation isn't needed there.
     */
    public static void addOpensForApt() {
        Class<?> cModule;
        try {
            cModule = Class.forName("java.lang.Module");
        } catch (ClassNotFoundException e) {
            return; //jdk8-; this is not needed.
        }

        Object jdkCompilerModule = UnsafeUtil.getJdkCompilerModule();
        Object ownModule = getOwnModule();
        try {
            Method m = cModule.getDeclaredMethod("implAddOpens", String.class, cModule);
            UnsafeUtil.putBooleanVolatile(m, true);
            for (String p : UnsafeUtil.JAVAC_PKG_S) {
                m.invoke(jdkCompilerModule, p, ownModule);
            }
        } catch (Exception ignore) {
        }
    }

    public static long getFirstFieldOffset(Unsafe unsafe) {
        try {
            return unsafe.objectFieldOffset(Parent.class.getDeclaredField("first"));
        } catch (NoSuchFieldException e) {
            // can't happen.
            throw new RuntimeException(e);
        } catch (SecurityException e) {
            // can't happen
            throw new RuntimeException(e);
        }
    }

    private static Object getOwnModule() {
        try {
            Method m = Permit.getMethod(Class.class, "getModule");
            return m.invoke(BaseProcessor.class);
        } catch (Exception e) {
            return null;
        }
    }

    @SuppressWarnings("unchecked")
    private static <T> T jbUnwrap(T wrapper) {
        addOpensForApt();
        if (wrapper instanceof JavacProcessingEnvironment env) {
            return (T) env;
        }

        T unwrapped = null;
        try {
            // 由于本项目需要在 IDEA 中进行调试和运行，因此引入 IDEA 环境的 ProcessingEnvironment。
            final Class<?> apiWrappers = wrapper.getClass().getClassLoader().loadClass("org.jetbrains.jps.javac.APIWrappers");
            final Method unwrapMethod = apiWrappers.getDeclaredMethod("unwrap", Class.class, Object.class);
            unwrapped = ((Class<? extends T>) ProcessingEnvironment.class)
                    .cast(unwrapMethod.invoke(null, ProcessingEnvironment.class, wrapper));
        } catch (Throwable ignored) {
        }
        return unwrapped != null ? unwrapped : wrapper;
    }


    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        processingEnv = jbUnwrap(processingEnv);
        super.init(processingEnv);

        this.context = ((JavacProcessingEnvironment) processingEnv).getContext();

        this.messager = processingEnv.getMessager();
        this.trees = JavacTrees.instance(processingEnv);
        this.maker = TreeMaker.instance(context);
        this.names = Names.instance(context);

        this.types = Types.instance(context);
        this.filer = processingEnv.getFiler();
        this.elements = processingEnv.getElementUtils();

        info("========= {} init =========", getClass().getSimpleName());
    }

    /**
     * 从环境变量获取包扫描路径，并扫描AbstractHandler实现装载到MAP中
     * 这么做的目的是为了处理外部扩展(自定义并且不在组件包路径下)的AbstractHandler
     * 但是当前并没有开放这个实现
     *
     * @see AbstractHandler 的内部扫包机制
     */
    private void register(ProcessingEnvironment env) {
        String handlerPackage = env.getOptions().get("apt.package");
        if (null == handlerPackage) {
            handlerPackage = APT_HANDLER_PACKAGE;
        }
        // 扫包
        Reflections reflections = new Reflections(new ConfigurationBuilder()
                .forPackages(handlerPackage)
                .addScanners(Scanners.SubTypes)
        );

        Set<Class<? extends AbstractHandler>> set = reflections.getSubTypesOf(AbstractHandler.class);
        for (Class<? extends AbstractHandler> subType : set) {
            try {
                AbstractHandler handler = subType.getDeclaredConstructor((Class<?>) null)
                        .newInstance();
                Set<String> handleTypes = handler.getHandleTypes();
                for (String handleType : handleTypes) {
                    MAP.put(handleType, handler);
                }
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
    }

    @SuppressWarnings("unchecked")
    public List<JCTree.JCAnnotation> getAnnotations(Element element) {
        Tree tree = trees.getTree(element);
        switch (tree.getKind()) {
            case VARIABLE -> {
                VariableTree variableTree = (VariableTree) tree;
                return (List<JCTree.JCAnnotation>) variableTree.getModifiers().getAnnotations();
            }
            case METHOD -> {
                MethodTree methodTree = (MethodTree) tree;
                return (List<JCTree.JCAnnotation>) methodTree.getModifiers().getAnnotations();
            }
            default -> {
                return null;
            }
        }
    }

     public Set<Modifier> getFlags(Element element) {
        Tree tree = trees.getTree(element);
         return switch (tree.getKind()) {
             case VARIABLE -> {
                 VariableTree variableTree = (VariableTree) tree;
                 yield variableTree.getModifiers().getFlags();
             }
             case METHOD -> {
                 MethodTree methodTree = (MethodTree) tree;
                 yield methodTree.getModifiers().getFlags();
             }
             default -> null;
         };
    }


    private List<JCTree.JCAnnotation> cloneAnnotations(Element element) {
        return List.from(element.getAnnotationMirrors()
                .stream()
                .map(annotationMirror -> {
                            java.util.List<JCTree.JCAssign> collect = annotationMirror.getElementValues().entrySet().stream()
                                    .map(entry -> maker.Assign(
                                            maker.Ident(names.fromString(entry.getKey().getSimpleName().toString())),
                                            maker.Literal(entry.getValue().getValue().toString())
                                    ))
                                    .collect(Collectors.toList());

                            String name = annotationMirror.getAnnotationType().asElement().getSimpleName().toString();
                            JCTree.JCIdent ident = maker.Ident(names.fromString(name));
                            return maker.Annotation(ident, List.from(collect));
                        }
                )
                .toArray(JCTree.JCAnnotation[]::new)
        );
    }


    public JCTree.JCClassDecl getClassDeclFromVariable(Element element) {
        ElementKind kind = element.getKind();
        if (ElementKind.CLASS == kind) {
            return trees.getTree((TypeElement) element);
        }

        Element enclosingElement = element.getEnclosingElement();
        if (ElementKind.CLASS == enclosingElement.getKind()) {
            return trees.getTree((TypeElement) enclosingElement);
        }

        throw new RuntimeException("Unknown type.");
    }


    public JCTree.JCAnnotation makeAnnotation(String annotationName, List<JCTree.JCExpression> args) {
        JCTree.JCExpression expression = chainDots(annotationName.split(Consts.PACKAGE_DELIMITER));
        return maker.Annotation(expression, args);
    }

    private void createNewField(VariableElement fieldElement) {
        // 通过字段获取类名
        String className = ((TypeElement) fieldElement.getEnclosingElement()).getQualifiedName().toString();
        String fieldName = fieldElement.getSimpleName().toString();
        TypeMirror fieldType = fieldElement.asType();
    }

    public JCTree.JCExpression makeArg(String key, Object value) {
        JCTree.JCExpression jcValue = AptUtil.toJcExpression(maker, names, value);
        return maker.Assign(maker.Ident(names.fromString(key)), jcValue);
    }

    public JCTree.JCExpression makeArgs(String key, Object... values) {
        ListBuffer<JCTree.JCExpression> list = new ListBuffer<>();
        for (Object value : values) {
            JCTree.JCExpression jcValue = AptUtil.toJcExpression(maker, names, value);
            list.append(jcValue);
        }

        // JCNewArray jcNewArray = maker.NewArray(maker.Ident(names.fromString("java.lang.String")), List.nil(),  list.toList());
        // 上面这一行这么写生成的就是 map = new String[]{"1:A", "2:B", "3:C"};

        // maker.NewArray不给Type参数，否则会报：java: 注释中不允许使用 'new'
        JCTree.JCNewArray newArray = maker.NewArray(null, List.nil(), list.toList());
        return maker.Assign(maker.Ident(names.fromString(key)), newArray);
    }

    protected JCTree.JCFieldAccess makeThisVarRef(Name varName) {
        return maker.Select(maker.Ident(makeName("this")), varName);
    }

    protected  Name makeName(String s) {
        return names.fromString(s);
    }

    protected JCTree.JCIdent makeIdent(String name) {
        return maker.Ident(makeName(name));
    }

    private JCTree.JCExpression chainDots(String... elems) {
        assert elems != null;

        JCTree.JCExpression e = null;
        for (String elem : elems) {
            e = e == null ? maker.Ident(names.fromString(elem)) : maker.Select(e, names.fromString(elem));
        }
        assert e != null;

        return e;
    }

    private JCTree.JCAssign assignClone(JCTree.JCAssign jcAssign) {
        JCTree.JCIdent left = maker.Ident(names.fromString(jcAssign.lhs.toString()));
        JCTree.JCIdent right = maker.Ident(names.fromString(jcAssign.rhs.toString()));
        return maker.Assign(left, right);
    }

    private JCTree.JCExpression argClone(JCTree.JCAssign jcAssign) {
        return makeArg(jcAssign.lhs.toString(), jcAssign.rhs.toString().replace("\"", ""));
    }

    private JCTree.JCExpression select(String components) {
        String[] componentArray = components.split("\\.");
        JCTree.JCExpression expr = maker.Ident(names.fromString(componentArray[0]));
        for (int i = 1; i < componentArray.length; i++) {
            expr = maker.Select(expr, names.fromString(componentArray[i]));
        }
        return expr;
    }


    public void info(String ss, Object... args) {
        log(Diagnostic.Kind.NOTE, ss, args);
    }

    public void warn(String ss, Object... args) {
        log(Diagnostic.Kind.WARNING, ss, args);
    }

    public void error(String ss, Object... args) {
        log(Diagnostic.Kind.ERROR, ss, args);
    }

    public void log(Diagnostic.Kind kind, String ss, Object... args) {
        if (args.length > 0) {
            ss = ss.replace("{}", "%s");
            String logs = String.format(ss, args);
            messager.printMessage(kind, logs);
        } else {
            messager.printMessage(kind, ss);
        }
    }

    public Context getContext() {
        return context;
    }

    public TreeMaker getMaker() {
        return maker;
    }

    public Names getNames() {
        return names;
    }

}